
#include "eye_tiling.h"
#include "register/op_def_registry.h"


namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  EyeTilingData tiling;
  const gert::StorageShape* y_shape = context->GetInputShape(0);

    int dim = y_shape->GetStorageShape().GetDimNum();
    int num_rows = y_shape->GetStorageShape().GetDim(dim-2);
    int num_columns = y_shape->GetStorageShape().GetDim(dim-1);
    int32_t batch_shape_size = 1;
    for (int i = 0; i < (dim - 2); i++)
    {
        batch_shape_size *= y_shape->GetStorageShape().GetDim(i);
    }
    int8_t formerNum;
    int32_t tailLength;
    {
        if(batch_shape_size < 32)
        {
            formerNum = batch_shape_size;
            tailLength = 0;
            context->SetBlockDim(formerNum);
        }
        else
        {
            formerNum = batch_shape_size % 32;
            tailLength = batch_shape_size / 32;
            if(formerNum == 0) 
            {
                formerNum = 32;
                tailLength -= 1;
            }
            context->SetBlockDim(32);
        }
    }

    int32_t num_rows_columns = num_rows * num_columns;
    if(num_rows > num_columns) num_rows = num_columns;
    tiling.set_num_rows(num_rows);
    tiling.set_num_columns(num_columns);
    tiling.set_num_rows_columns(num_rows_columns);
    tiling.set_formerNum(formerNum);
    tiling.set_tailLength(tailLength);

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Eye : public OpDef {
public:
    explicit Eye(const char* name) : OpDef(name)
    {
        this->Input("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_DOUBLE})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("num_rows").Int();
        this->Attr("num_columns").Int();
        this->Attr("batch_shape").ListInt({});
        this->Attr("dtype").Int();

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(Eye);
}
